package com.example.demo.util;

import brave.Span;
import brave.Tracer;
import brave.propagation.TraceContext;
import lombok.extern.slf4j.Slf4j;
import org.slf4j.MDC;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import javax.annotation.PostConstruct;

/**
 * TraceUtil
 *
 * @author wangchen
 * @date 2024/6/9
 **/
@Slf4j
@Component
public class TraceUtil {

    // 静态依赖变量（由Spring注入后赋值）
    private static Tracer tracer;

    // 实例变量（用于接收Spring注入）
    @Autowired
    private Tracer instanceTracer;

    // 初始化标记（确保依赖注入完成）
    private static boolean initialized = false;

    @PostConstruct
    public void init() {
        try {
            // 将实例变量赋值给静态变量
            tracer = this.instanceTracer;
            initialized = true;
            log.info("TraceUtil初始化完成，静态方法可调用");
        } catch (Exception e) {
            log.error("TraceUtil初始化失败", e);
            throw new RuntimeException("TraceUtil初始化失败，请检查Spring容器依赖是否正确配置", e);
        }
    }

    /**
     * 校验工具类是否已初始化
     */
    private static void checkInitialized() {
        if (!initialized) {
            throw new IllegalStateException("TraceUtil尚未初始化，请检查Spring配置是否正确注入依赖");
        }
    }

    /**
     * 将当前 TraceId 放入 MDC 上下文（供日志框架使用）
     */
    public static void putTraceIdToMdc() {
        TraceContext context = getContext();
        if (context == null) {
            return;
        }
        // "traceId" 为日志中引用的键
        MDC.put("X-B3-TraceId", context.traceIdString());
        MDC.put("X-B3-SpanId", context.spanIdString());
        MDC.put("X-B3-ParentSpanId", context.parentIdString());
    }

    /**
     * 从 MDC 中清除 TraceId（避免线程复用导致的信息污染）
     */
    public static void removeTraceIdFromMdc() {
        MDC.remove("X-B3-TraceId");
        MDC.remove("X-B3-SpanId");
        MDC.remove("X-B3-ParentSpanId");
    }

    /**
     * 获取traceId
     */
    public static String getTraceId() {
        TraceContext context = getContext();
        if (context == null) {
            return null;
        }
        try {
            return context.traceIdString();
        } catch (Exception e) {
            log.error("获取TraceId异常", e);
            return null;
        }
    }

    /**
     * 获取context
     */
    private static TraceContext getContext() {
        checkInitialized();
        Span span = tracer.currentSpan();
        if (span == null) {
            span = tracer.nextSpan();
            span.start();
        }
        try {
            return span.context();
        } catch (Exception e) {
            log.error("获取TraceContext异常", e);
            return null;
        }
    }
}
